# visual_pred.py

import os
import os.path as osp
import sys
import glob
import logging
import argparse
import numpy as np
from tqdm import tqdm
import cv2
import torch
import re
from pathlib import Path
from scipy import ndimage

# --- Import from our visual utility file ---
from visual_util import launch_visualizer


def setup_logging(log_file='verification.log'):
    """Configures logging to write to a file and the console."""
    logger = logging.getLogger('VerificationLogger')
    logger.setLevel(logging.INFO)
    if logger.hasHandlers():
        logger.handlers.clear()
    file_handler = logging.FileHandler(log_file, mode='w')
    file_formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
    file_handler.setFormatter(file_formatter)
    logger.addHandler(file_handler)
    console_handler = logging.StreamHandler()
    console_formatter = logging.Formatter('%(message)s')
    console_handler.setFormatter(console_formatter)
    logger.addHandler(console_handler)
    return logger

def imread_cv2(path, options=cv2.IMREAD_COLOR):
    """Open an image or a depthmap with opencv-python."""
    if path.endswith((".exr", "EXR")):
        options = cv2.IMREAD_ANYDEPTH
    img = cv2.imread(path, options)
    if img is None:
        raise IOError(f"Could not load image={path} with {options=}")
    if img.ndim == 3:
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    return img

def detect_sky_mask(img_rgb):
    """
    Detect sky pixels using HSV color space and morphological operations.
    Args:
        img_rgb: RGB image normalized to [-1, 1]
    Returns:
        Boolean mask (as int8) where True indicates non-sky pixels.
    """
    img = ((img_rgb + 1) * 127.5).astype(np.uint8)
    img_bgr = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
    hsv = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2HSV)

    lower_blue = np.array([105, 50, 140])
    upper_blue = np.array([135, 255, 255])
    mask_blue = cv2.inRange(hsv, lower_blue, upper_blue)

    lower_light_blue = np.array([95, 5, 150])
    upper_light_blue = np.array([145, 100, 255])
    mask_light_blue = cv2.inRange(hsv, lower_light_blue, upper_light_blue)

    lower_white = np.array([0, 0, 235])
    upper_white = np.array([180, 10, 255])
    mask_white = cv2.inRange(hsv, lower_white, upper_white)

    mask = mask_blue | mask_light_blue | mask_white

    height = mask.shape[0]
    upper_third = int(height * 0.4)
    upper_region = hsv[:upper_third, :, :]
    mask[:upper_third, :] |= ((upper_region[:, :, 1] < 50) & (upper_region[:, :, 2] > 150))

    kernel = np.ones((7, 7), np.uint8)
    mask = cv2.dilate(mask, kernel, iterations=1)
    mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)

    mask = mask.astype(bool)
    labels, num_labels = ndimage.label(mask)
    if num_labels > 0:
        top_row_labels = set(labels[0, :])
        top_row_labels.discard(0)
        if top_row_labels:
            mask = np.isin(labels, list(top_row_labels))
            labels, num_labels = ndimage.label(mask)
            if num_labels > 0:
                sizes = ndimage.sum(mask, labels, range(1, num_labels + 1))
                mask_size = mask.size
                big_enough = sizes > mask_size * 0.01
                mask = np.isin(labels, np.where(big_enough)[0] + 1)
    return (~mask).astype(np.int8)

def is_outdoor_scene(mask):
    sky_ratio = 1.0 - np.mean(mask)
    significant = sky_ratio > 0.2
    return significant


def load_prediction_data(view_path, logger, gui_min_conf_percentile=10.0, local=False, other=True, gt=False,pcent_filter=True):
    """
    Loads all necessary data for a single predicted view.
    This includes the predicted point cloud, the input image, and the camera parameters.
    """
    try:
        base_path = view_path.removesuffix(".png").removesuffix(".jpg").removesuffix("_rgb")
        pred_path = f"{base_path}.npz"

        if not all(osp.exists(p) for p in [pred_path]):
            logger.warning(f"Missing files for {osp.basename(base_path)}. Skipping.")
            return None

        rgb_image = imread_cv2(view_path, cv2.IMREAD_COLOR).astype(np.float32) / 255.0
        if rgb_image is None:
            logger.warning(f"Could not read image {view_path}. Skipping.")
            return None

        pred = np.load(pred_path, allow_pickle=True)
        if 'intrinsics_gt' in pred:
            intrinsics = pred["intrinsics_gt"].astype(np.float32)
        elif pred["intrinsics"] == None:
            intrinsics = None
        else:
            intrinsics = pred["intrinsics"].astype(np.float32)
        camera_pose = pred["pose"].astype(np.float32)

        not_sky_mask = detect_sky_mask(rgb_image).flatten().astype(np.int8)

        pts3d_global = pred['pts3d_in_other_view'].reshape(-1, 3)
        conf_global = pred['conf'].flatten()
        if "pts3d_in_self_view_aligned_to_global" in pred:
            pts3d_local = pred['pts3d_in_self_view_aligned_to_global'].reshape(-1, 3)
        else: 
            pts3d_local = pred['pts3d_local_aligned_to_global'].reshape(-1, 3)
        if "conf_self" in pred:
            conf_local = pred['conf_self'].flatten()
        else:
            conf_local = pred['conf_local'].flatten()
        img_rgb_flat = rgb_image.reshape(-1, 3)

        # if is_outdoor_scene(not_sky_mask):
        #     pts3d_global = pts3d_global[not_sky_mask > 0]
        #     pts3d_local = pts3d_local[not_sky_mask > 0]
        #     img_rgb_flat = img_rgb_flat[not_sky_mask > 0]
        #     conf_global = conf_global[not_sky_mask > 0]
        #     conf_local = conf_local[not_sky_mask > 0]
        
        if pcent_filter:
            sort_idx_global = np.argsort(-conf_global)
            sorted_conf_global = conf_global[sort_idx_global]
            sorted_pts3d_global = pts3d_global[sort_idx_global]
            sorted_img_rgb_global = img_rgb_flat[sort_idx_global]
            # sorted_not_sky_global = not_sky_mask[sort_idx_global]

            sort_idx_local = np.argsort(-conf_local)
            sorted_conf_local = conf_local[sort_idx_local]
            sorted_pts3d_local = pts3d_local[sort_idx_local]
            sorted_img_rgb_local = img_rgb_flat[sort_idx_local]
            # sorted_not_sky_local = not_sky_mask[sort_idx_local]        
        
            total_global = len(sorted_pts3d_global)
            total_local = len(sorted_pts3d_local)

            # Calculate number of points to show based on the percentile
            num_global = max(1, int(total_global * (100 - gui_min_conf_percentile) / 100))
            num_local = max(1, int(total_local * (100 - gui_min_conf_percentile) / 100))

            pts3d_global = sorted_pts3d_global[:num_global]
            pts3d_local = sorted_pts3d_local[:num_local]
            img_rgb_global = sorted_img_rgb_global[:num_global]
            img_rgb_local = sorted_img_rgb_local[:num_local] 
            conf_global = sorted_conf_global[:num_global]
            conf_local = sorted_conf_local[:num_local]
        else:
            # print(conf_global.max(), conf_global.min(),conf_global.mean(), conf_local.max(), conf_local.min(),conf_local.mean())
            # print((conf_global>gui_min_conf_percentile).all(), (conf_local>gui_min_conf_percentile).all())
            # print((conf_global>gui_min_conf_percentile).sum(), (conf_local>gui_min_conf_percentile).sum(),(conf_global<=gui_min_conf_percentile).sum(), (conf_local<=gui_min_conf_percentile).sum())
            pts3d_global = pts3d_global[conf_global>gui_min_conf_percentile]
            img_rgb_global = img_rgb_flat[conf_global>gui_min_conf_percentile]
            conf_global = conf_global[conf_global>gui_min_conf_percentile]
            pts3d_local = pts3d_local[conf_local>gui_min_conf_percentile]
            img_rgb_local = img_rgb_flat[conf_local>gui_min_conf_percentile] 
            conf_local = conf_local[conf_local>gui_min_conf_percentile]
                
        
        # img_rgb_global = img_rgb_local = img_rgb_flat
        results =[]
        if local:
            results.append({
                    "rgb": img_rgb_local,
                    "pts": pts3d_local,
                    "conf": conf_local,
                    "intrinsics": intrinsics,
                    "pose": camera_pose,
                    "view_name": f"{osp.basename(base_path)}_local"
                })
        if other:
            results.append({
                    "rgb": img_rgb_global,
                    "pts": pts3d_global,
                    "conf": conf_global,
                    "intrinsics": intrinsics,
                    "pose": camera_pose,
                    "view_name": osp.basename(base_path)
                })
        if gt:
            results.append({
                    "rgb": rgb_image.reshape(-1, 3),
                    "pts": pred["pts3d_gt"].reshape(-1, 3),
                    "conf": np.ones(pred["pts3d_gt"].reshape(-1, 3).shape[0]),
                    "intrinsics": intrinsics,
                    "pose": camera_pose,
                    "view_name": f"{osp.basename(base_path)}_gt"
                })
        return results

    except Exception as e:
        logger.error(f"Could not process view for {osp.basename(view_path)}. Details: {e}", exc_info=True)
        return None

# ======================================================================================
# --- Main Execution Block ---
# ======================================================================================
def view_folder(sample_dir, args):
    sample_dir = Path(sample_dir)
    sample_key = f"{sample_dir.parent.name}/{sample_dir.name}"
    logger.info(f"\n--- Processing sample: {sample_key} ---")

    rgb_paths = list(sample_dir.glob("*.png")) + list(sample_dir.glob("*.jpg"))
    if not rgb_paths:
        logger.warning(f"No RGB images found in {sample_dir}. Skipping.")
        return False

    # --- Prepare data for the visualizer ---
    pc_list, color_list, conf_list, view_names = [], [], [], []
    focals, pps, Rs, ts = [], [], [], []

    logger.info("Loading views and predicted point clouds...")
    for rgb_path in tqdm(rgb_paths, file=sys.stdout):
        view_data_list = load_prediction_data(
            str(rgb_path),
            logger,
            gui_min_conf_percentile=2.1,
            local=args.local, other=args.other, gt=args.gt,
            pcent_filter=False
        )
        for view_data in view_data_list:
            pc_list.append(view_data["pts"])
            color_list.append(view_data['rgb'])
            view_names.append(osp.basename(view_data["view_name"]))
            
            # Create a confidence map. Here we assume all predicted points are valid.
            # The shape should correspond to the image shape.
            conf_list.append(view_data['conf'])

            # Aggregate camera parameters for the viewer
            intrinsics = view_data['intrinsics']
            pose = view_data['pose']
            if intrinsics is not None:
                focals.append(intrinsics[0, 0])
                pps.append((intrinsics[0, 2], intrinsics[1, 2]))
            Rs.append(pose[:3, :3])
            ts.append(pose[:3, 3])

    if not pc_list:
        logger.warning(f"No valid data could be loaded for {sample_key}. Skipping visualization.")
        return True

    cam_dict = {'focal': focals, 'pp': pps, 'R': Rs, 't': ts}

    # --- Launch the visualizer ---
    try:
        logger.info(f"Launching interactive visualizer for {sample_key} at http://localhost:{args.port}")
        logger.info("--> Close the browser tab or press Ctrl+C in the console to continue to the next sample.")
        launch_visualizer(
            pc_list=pc_list,
            color_list=color_list,
            conf_list=conf_list,
            cam_dict=cam_dict,
            view_names=view_names,
            port=args.port
        )
    except KeyboardInterrupt:
        logger.info(f"\nVisualizer for {sample_key} closed, proceeding...")
    except Exception as e:
        logger.error(f"Failed to launch visualizer for {sample_key}: {e}", exc_info=True)

    return True

if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description="Visually verify predicted point clouds."
    )
    
    parser.add_argument(
        "--pred_root", 
        type=str,
        default=r"D:\data\3D\eval_outdir\examples\room\images\0",
        help="Path to the root directory of the predictions."
    )

    parser.add_argument(
        "--log_file", 
        type=str, 
        default="logs/visual_pred.log",
        help="Path to the output log file."
    )

    parser.add_argument(
        "--port", 
        type=int, 
        default=8080,
        help="Port to run the Viser web server on."
    )

    parser.add_argument(
        '--local',
        action='store_true',
        help="Load and visualize the local point cloud."
    )
    parser.add_argument(
        '--other',
        action='store_true',
        help="Load and visualize the global (other) point cloud."
    )
    parser.add_argument(
        '--gt',
        action='store_true',
        help="Load and visualize the ground truth point cloud."
    )

    
    args = parser.parse_args()
    
    # Create logs directory if it doesn't exist
    os.makedirs(osp.dirname(args.log_file), exist_ok=True)
    logger = setup_logging(args.log_file)

    if not osp.isdir(args.pred_root):
        logger.error(f"Provided prediction root '{args.pred_root}' does not exist.")
        sys.exit(1)
    
    # If no specific point cloud type is requested, default to showing all.
    if not (args.local or args.other or args.gt):
        logger.info("No specific point cloud type requested. Defaulting to show all (--local --other --gt).")
        args.local = args.other = args.gt = True

    view_folder(args.pred_root, args)

    logger.info("\n--- Visual prediction verification complete. ---")
